#!/bin/sh -e

EFIDIR=$(grep ^ID= /etc/os-release | sed -e 's/^ID=//' -e 's/rhel/redhat/' -e 's/\"//g')
if [ -d /sys/firmware/efi/efivars/ ]; then
    grubdir=`echo "/@bootdirname@/efi/EFI/${EFIDIR}/" | sed 's,//*,/,g'`
else
    grubdir=`echo "/@bootdirname@/@grubdirname@" | sed 's,//*,/,g'`
fi

PACKAGE_VERSION="@PACKAGE_VERSION@"
PACKAGE_NAME="@PACKAGE_NAME@"
self=`basename $0`
bindir="@bindir@"
grub_mkpasswd="${bindir}/@grub_mkpasswd_pbkdf2@"

# Usage: usage
# Print the usage.
usage () {
    cat <<EOF
Usage: $0 [OPTION]
$0 prompts the user to set a password on the grub bootloader. The password
is written to a file named user.cfg which lives in the GRUB directory
located by default at ${grubdir}.

  -h, --help                     print this message and exit
  -v, --version                  print the version information and exit
  -o, --output_path <DIRECTORY>  put user.cfg in a user-selected directory

Report bugs at https://bugzilla.redhat.com.
EOF
}

argument () {
    opt=$1
    shift

    if test $# -eq 0; then
        gettext_printf "%s: option requires an argument -- \`%s'\n" "$self" "$opt" 1>&2
        exit 1
    fi
    echo $1
}

# Ensure that it's the root user running this script
if [ "${EUID}" -ne 0 ]; then
    echo "The grub bootloader password may only be set by root."
    usage
    exit 2
fi

# Check the arguments.
while test $# -gt 0
do
    option=$1
    shift

    case "$option" in
    -h | --help)
	usage
	exit 0 ;;
    -v | --version)
	echo "$self (${PACKAGE_NAME}) ${PACKAGE_VERSION}"
	exit 0 ;;
    -o | --output)
        OUTPUT_PATH=`argument $option "$@"`; shift ;;
    --output=*)
        OUTPUT_PATH=`echo "$option" | sed 's/--output=//'` ;;
    -o=*)
        OUTPUT_PATH=`echo "$option" | sed 's/-o=//'` ;;
    esac
done

# set user input or default path for user.cfg file
if [ -z "${OUTPUT_PATH}" ]; then
    OUTPUT_PATH="${grubdir}"
fi

if [ ! -d "${OUTPUT_PATH}" ]; then
    echo "${OUTPUT_PATH} does not exist."
    usage
    exit 2;
fi

ttyopt=$(stty -g)
fixtty() {
      stty ${ttyopt}
}

trap fixtty EXIT

getsaltpass() {
	local P0
	local P1
	P0="$1" && shift
	P1="$1" && shift
	P2="$1" && shift

	( echo ${P0} ; echo ${P1} ) | \
    	LC_ALL=C ${grub_mkpasswd} -a ${P2} | \
    	grep -v '[eE]nter password:' | \
    	sed -e "s/PBKDF2 hash of your password is //"
}

verifyusercfgoldpasswd() {
    # get old password salt
    expectsalt=`cat ${grubdir}/user.cfg | cut -d "." -f 5`
    # get expect password
    expectpass=`cat ${grubdir}/user.cfg`
    prefix="GRUB2_PASSWORD="

    stty -echo
    echo -n "Enter Current password: "
    read PASSWORD_CURRENT
    echo

    needcheckpass="${prefix}$(getsaltpass "${PASSWORD_CURRENT}" "${PASSWORD_CURRENT}" "${expectsalt}")"
    if [ "$expectpass" != "$needcheckpass" ]; then
        echo "Authentication failed"
        exit 1
    fi

    stty ${ttyopt}
}

verifygrubcfgoldpasswd() {
    # get old password line
    expectpass=`cat ${grubdir}/grub.cfg | grep "password_pbkdf2 root grub.pbkdf2.sha512" | cut -d " " -f 3`
    # if not get password, try a quotation mark match
    if [ -z "$expectpass" ];then
        expectpass=`cat ${grubdir}/grub.cfg | grep "password_pbkdf2 root \"grub.pbkdf2.sha512" | cut -d " " -f 3 | cut -d "\"" -f 2`
    fi
    if [ -z "$expectpass" ];then
        expectpass=`cat ${grubdir}/grub.cfg | grep "password_pbkdf2 root 'grub.pbkdf2.sha512" | cut -d " " -f 3 | cut -d "'" -f 2`
    fi
    if [ -n "$expectpass" ];then
        # get old password salt
        expectsalt=`echo ${expectpass} | cut -d "." -f 5`
        stty -echo
        echo -n "Enter Current password: "
        read PASSWORD_CURRENT
        echo

        needcheckpass="$(getsaltpass "${PASSWORD_CURRENT}" "${PASSWORD_CURRENT}" "${expectsalt}")"
        if [ "$expectpass" != "$needcheckpass" ]; then
            echo "Authentication failed"
            exit 1
        fi
    fi

}

if [ -e ${grubdir}/user.cfg ];then
    verifyusercfgoldpasswd
else
    verifygrubcfgoldpasswd
fi

checkcomplexity() {
    set +e
    USERNAME=`cat ${grubdir}/grub.cfg | grep "set superusers=" | cut -d "\"" -f 2 |tail -1`
    local P1="$1" && shift
    if [ "$P1" = "$USERNAME" ];then
        echo "The password contains the user name in some form"
        exit 1
    fi
    # password len >= 8
    strlen=`echo "$P1" | grep -E '^(.{8,}).*$'`
    if [ -z "$strlen" ];then
        echo "The password is shorter than 8 characters"
        exit 1
    fi
    # lowercase
    strlow=`echo "$P1" | grep -E --color '^(.*[a-z]+).*$'`
    # uppercase
    strupp=`echo $P1 | grep -E --color '^(.*[A-Z]).*$'`
    # special character
    strts=`echo $P1 | grep -E --color '^(.*\W).*$'`
    # num
    strnum=`echo $P1 | grep -E --color '^(.*[0-9]).*$'`
    complexity=0
    if [ -n "$strlow" ];then
        complexity=`expr $complexity + 1`
    fi
    if [ -n "$strupp" ];then
        complexity=`expr $complexity + 1`
    fi
    if [ -n "$strts" ];then
        complexity=`expr $complexity + 1`
    fi
    if [ -n "$strnum" ];then
        complexity=`expr $complexity + 1`
    fi
    if [ $complexity -lt 3 ];then
        echo "The password contains less than 3 character classes"
        exit 1
    fi
    set -e
}

stty -echo

# prompt & confirm new grub2 root user password
echo -n "Enter password: "
read PASSWORD
echo
stty ${ttyopt}
checkcomplexity $PASSWORD
stty -echo
echo -n "Confirm password: "
read PASSWORD_CONFIRM
echo
stty ${ttyopt}
checkcomplexity $PASSWORD_CONFIRM

getpass() {
    local P0
    local P1
    P0="$1" && shift
    P1="$1" && shift

    ( echo ${P0} ; echo ${P1} ) | \
        LC_ALL=C ${grub_mkpasswd} | \
        grep -v '[eE]nter password:' | \
        sed -e "s/PBKDF2 hash of your password is //"
}

MYPASS="$(getpass "${PASSWORD}" "${PASSWORD_CONFIRM}")"
if [ -z "${MYPASS}" ]; then
      echo "${self}: error: empty password" 1>&2
      exit 1
fi

# on the ESP, these will fail to set the permissions, but it's okay because
# the directory is protected.
install -m 0600 /dev/null "${OUTPUT_PATH}/user.cfg" 2>/dev/null || :
chmod 0600 "${OUTPUT_PATH}/user.cfg" 2>/dev/null || :
echo "GRUB2_PASSWORD=${MYPASS}" > "${OUTPUT_PATH}/user.cfg"

if ! grep -q "^### BEGIN /etc/grub.d/01_users ###$" "${OUTPUT_PATH}/grub.cfg"; then
    echo "WARNING: The current configuration lacks password support!"
    echo "Update your configuration with @grub_mkconfig@ to support this feature."
fi
